/*
 * MIT License
 *
 * Copyright (c) 2021 cao.yong<cymsn81@126.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package org.cy.websocket.stomp;

import org.cy.websocket.stomp.events.*;
import org.cy.websocket.stomp.messaging.GenericMessage;
import org.cy.websocket.stomp.messaging.Message;
import org.cy.websocket.stomp.messaging.StompMessage;
import org.cy.websocket.stomp.messaging.StompMessageConverter;
import org.cy.websocket.stomp.protocols.Versions;
import org.cy.websocket.stomp.timer.TimedTaskScheduler;
import org.cy.websocket.stomp.util.StringKit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * created in 2021/5/12
 *
 * @author cao.yong
 */
public class StompHandler {

    private static final Logger LOG = LoggerFactory.getLogger(StompHandler.class);

    private static final AtomicInteger COUNTER = new AtomicInteger(0);

    private static final ConcurrentMap<String, OnSubscribeEvent> ALL_SUBSCRIPTION = new ConcurrentHashMap<>(2);
    private static final ConcurrentMap<String, OnReceiptEvent> ALL_RECEIPT = new ConcurrentHashMap<>(1);

    private static final String TASK_PING = "PING";
    private static final String TASK_PONG = "PONG";
    private final TimedTaskScheduler scheduler;
    /**
     * Stomp 协议套接字包装类
     */
    private StompWebSocket socket;
    private String brokerUrl;
    private Heartbeat heartbeat;
    private long lastServerPing;
    private OnConnectedEvent onConnectedEvent;
    private OnDisconnectedEvent onDisconnectedEvent;
    private OnClosedEvent onClosedEvent;
    private OnErrorEvent onErrorEvent;

    private String connectVersion;
    private Constants.StompSocketState connectState;

    public StompHandler() {
        this.scheduler = new TimedTaskScheduler();
        this.heartbeat = new Heartbeat();
    }

    public Constants.StompSocketState getConnectState() {
        return connectState;
    }

    public String getConnectVersion() {
        return connectVersion;
    }

    public boolean isConnected() {
        return (Constants.StompSocketState.OPEN == this.getConnectState());
    }

    /**
     * 启动并连接到 STOMP 服务
     *
     * @param headers {@link StompHeaders} STOMP 协议头
     */
    public void start(StompHeaders headers) {
        this.socket.connect(this.brokerUrl, headers, new StompWebSocketListener() {
            @Override
            public void onOpen() {
                onWebSocketConnect(headers);
            }

            @Override
            public void onMessage(Message<byte[]> message) {
                onReceiveMessage(message);
            }

            @Override
            public void onClosed(int code, String reason) {
                onWebSocketClosed(code, reason);
            }

            @Override
            public void onFailure(Throwable ex) {
                onWebSocketError(ex);
            }
        });
    }

    /**
     * 发送消息到服务器
     *
     * @param message 消息
     */
    public void send(GenericMessage<byte[]> message) {
        this.socket.send(message);
    }

    /**
     * 发送消息到指定的目标通道
     *
     * @param destination 目标通道
     * @param headers     {@link StompHeaders} STOMP 协议头
     * @param content     发送内容
     */
    public void send(String destination, StompHeaders headers, String content) {
        byte[] payload = content == null ? new byte[0] : content.getBytes(StandardCharsets.UTF_8);
        headers.setDestination(destination);
        GenericMessage<byte[]> message = new GenericMessage<>(StompCommand.SEND, headers, payload);
        this.send(message);
    }

    /**
     * 订阅目标通道
     *
     * @param destination 目标通道
     * @param headers     {@link StompHeaders} STOMP 协议头
     * @param callback    订阅后接收通道消息的回调函数
     * @return 订阅实例，可用于取消订阅
     */
    public StompSubscription subscribe(String destination, StompHeaders headers, OnSubscribeEvent callback) {
        if (null == headers) {
            headers = new StompHeaders();
        }
        if (null == headers.getId() || headers.getId().length() == 0) {
            headers.setId("SUB-" + COUNTER.incrementAndGet());
        }
        headers.setDestination(destination);
        ALL_SUBSCRIPTION.putIfAbsent(headers.getId(), callback);
        this.send(new GenericMessage<>(StompCommand.SUBSCRIBE, headers, new byte[0]));
        final String id = headers.getId();
        return new StompSubscription(id, inHeaders -> unsubscribe(id, inHeaders));
    }

    /**
     * 取消订阅
     *
     * @param id      订阅的标识
     * @param headers {@link StompHeaders} STOMP 协议头
     */
    public void unsubscribe(String id, StompHeaders headers) {
        if (null == headers) {
            headers = new StompHeaders();
        }
        headers.setId(id);
        ALL_SUBSCRIPTION.remove(id);
        this.send(new GenericMessage<>(StompCommand.UNSUBSCRIBE, headers, new byte[0]));
    }

    /**
     * 发送断开消息到 STOMP 服务
     *
     * @param headers {@link StompHeaders} STOMP 协议头
     */
    public void disconnect(StompHeaders headers) {
        if (StringKit.isEmpty(headers.getReceipt())) {
            headers.setReceipt("ClientOnClose-" + COUNTER.incrementAndGet());
        }
        // 安全断开，获取服务器确认收据
        ALL_RECEIPT.put(headers.getReceipt(), message -> {
            closeWebSocket();
        });
        this.send(new GenericMessage<>(StompCommand.DISCONNECT, headers, new byte[0]));
    }

    protected void setSocket(StompWebSocket socket) {
        this.socket = socket;
    }

    protected void setBrokerUrl(String brokerUrl) {
        this.brokerUrl = brokerUrl;
    }

    protected void setHeartbeat(Heartbeat heartbeat) {
        this.heartbeat = heartbeat;
    }

    protected void setOnConnectedEvent(OnConnectedEvent onConnectedEvent) {
        this.onConnectedEvent = onConnectedEvent;
    }

    protected void setOnDisconnectedEvent(OnDisconnectedEvent onDisconnectedEvent) {
        this.onDisconnectedEvent = onDisconnectedEvent;
    }

    protected void setOnClosedEvent(OnClosedEvent onClosedEvent) {
        this.onClosedEvent = onClosedEvent;
    }

    protected void setOnErrorEvent(OnErrorEvent onErrorEvent) {
        this.onErrorEvent = onErrorEvent;
    }

    protected void commit(String transactionId) {
        StompHeaders headers = new StompHeaders();
        headers.set("transaction", transactionId);
        this.send(new GenericMessage<>(StompCommand.COMMIT, headers, new byte[0]));
    }

    protected void abort(String transactionId) {
        StompHeaders headers = new StompHeaders();
        headers.set("transaction", transactionId);
        this.send(new GenericMessage<>(StompCommand.ABORT, headers, new byte[0]));
    }

    public StompTransaction begin(String transactionId) {
        final String txId = StringKit.isNotEmpty(transactionId) ? transactionId : "tx-" + COUNTER.incrementAndGet();
        StompHeaders headers = new StompHeaders();
        headers.set("transaction", txId);
        this.send(new GenericMessage<>(StompCommand.BEGIN, headers, new byte[0]));
        return new StompTransaction() {
            @Override
            public String getId() {
                return txId;
            }

            @Override
            public void commit() {
                StompHandler.this.commit(txId);
            }

            @Override
            public void abort() {
                StompHandler.this.abort(txId);
            }
        };
    }

    public void ack(String messageId, String subscriptionId, StompHeaders headers) {
        if (null == headers) {
            headers = new StompHeaders();
        }
        if (Versions.V1_2.equals(getConnectVersion())) {
            headers.setId(messageId);
        } else {
            headers.setMessageId(messageId);
        }
        headers.setSubscription(subscriptionId);
        this.send(new GenericMessage<>(StompCommand.ACK, headers, new byte[0]));
    }

    public void nack(String messageId, String subscriptionId, StompHeaders headers) {
        if (null == headers) {
            headers = new StompHeaders();
        }
        if (Versions.V1_2.equals(getConnectVersion())) {
            headers.setId(messageId);
        } else {
            headers.setMessageId(messageId);
        }
        headers.setSubscription(subscriptionId);
        this.send(new GenericMessage<>(StompCommand.NACK, headers, new byte[0]));
    }

    /**
     * 设置心跳
     *
     * @param headers 连接头信息
     */
    private void heartbeatSetting(StompHeaders headers) {
        if (Versions.V1_1.equals(getConnectVersion()) || Versions.V1_2.equals(getConnectVersion())) {
            /*
             * 客户端 PING
             * 判断客户端是否向服务器端发送 PING 消息，同时判断服务器端是否接收客户端 PONG 消息
             * 任何一方设置为 0 客户端都不会发送 PING 消息
             */
            long serverIncoming = headers.getHeartbeat()[1];
            if (this.heartbeat.getOutgoing() > 0 && serverIncoming > 0) {
                long ttl = Math.max(this.heartbeat.getOutgoing(), serverIncoming);
                LOG.debug("Send PING every {}ms", ttl);
                this.scheduler.start(TASK_PING, () -> {
                    send(new GenericMessage<>(StompCommand.HEARTBEAT, headers, new byte[]{0x0A}));
                    LOG.debug(">>> PING");
                }, ttl);
            }
            /*
             * 客户端 PONG
             *
             * 判断服务器端是否向客户端发送 PING 消息，同时判断客户端是否接收服务器端 PONG 消息
             * 任何一方设置为 0 客户端都不会处理 PONG 消息
             */
            long serverOutgoing = headers.getHeartbeat()[0];
            if (this.heartbeat.getIncoming() > 0 && serverOutgoing > 0) {
                long ttl = Math.max(this.heartbeat.getIncoming(), serverOutgoing);
                LOG.debug("Recv PONG every {}ms", ttl);
                this.scheduler.start(TASK_PONG, () -> {
                    // 这里监听服务器心跳消息最后一次时间
                    long interval = System.currentTimeMillis() - lastServerPing;
                    LOG.debug("间隔时间: {}ms", interval);
                    int rate = 2; // 倍比
                    if (interval > ttl * rate) {
                        disconnect(headers);
                    }
                }, ttl);
            }
        }
    }

    private OnSubscribeEvent getSubscribeEvent(String id) {
        return ALL_SUBSCRIPTION.get(id);
    }

    /**
     * 当 WebSocket 通道已打开
     *
     * @param headers {@link StompHeaders} STOMP 协议头
     */
    private void onWebSocketConnect(StompHeaders headers) {
        LOG.debug("Web Socket Opened...");
        this.connectState = Constants.StompSocketState.CONNECTING;
        headers.setAcceptVersion(Versions.supportedVersions());
        headers.setHeartbeat(heartbeat.getOutgoing(), heartbeat.getIncoming());
        GenericMessage<byte[]> message = new GenericMessage<>(StompCommand.CONNECT, headers, new byte[0]);
        this.send(message);
    }

    /**
     * 接收消息处理
     *
     * @param message 消息
     */
    private void onReceiveMessage(Message<byte[]> message) {
        this.lastServerPing = System.currentTimeMillis();
        switch (message.getCommand()) {
            case CONNECTED:
                connectState = Constants.StompSocketState.OPEN;
                connectVersion = message.getHeaders().getVersion();
                // 设置心跳
                heartbeatSetting(message.getHeaders());
                onConnectedEvent.on(message);
                break;
            case MESSAGE:
                String subscription = message.getHeaders().getSubscription();
                OnSubscribeEvent onReceive = getSubscribeEvent(subscription);
                StompMessage<byte[]> stompMessage = StompMessageConverter.getStompMessage(this, message);
                if (null != onReceive) {
                    onReceive.on(stompMessage);
                }
                break;
            case RECEIPT:
                OnReceiptEvent onReceipt = ALL_RECEIPT.get(message.getHeaders().getReceiptId());
                if (null != onReceipt) {
                    onReceipt.on(message);
                }
                break;
            case HEARTBEAT:
                LOG.debug("<<< PONG");
                break;
            default:
                LOG.debug("UNKNOWN: {} : {} : [{}]", message.getCommand(), message.getHeaders().toString(), message.getPayload().length);
                break;
        }
    }

    /**
     * WebSocket 通道已关闭
     *
     * @param code   关闭代码
     * @param reason 关闭原因
     */
    private void onWebSocketClosed(int code, String reason) {
        LOG.debug("Web Socket Closed: {}:{}", code, reason);
        this.connectState = Constants.StompSocketState.CLOSED;
        this.cleanUp();
        this.onDisconnectedEvent.on();
        this.onClosedEvent.on(code, reason);
    }

    /**
     * WebSocket 通道发生错误
     *
     * @param ex 异常错误实例
     */
    private void onWebSocketError(Throwable ex) {
        LOG.debug("Web Socket Error: {}", ex.getMessage());
        connectState = Constants.StompSocketState.CLOSED;
        this.onErrorEvent.on(ex);
    }

    private void closeWebSocket() {
        this.socket.close();
    }

    private void cleanUp() {
        this.scheduler.stop(TASK_PING);
        this.scheduler.stop(TASK_PONG);
    }
}
